import os 
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import time
import torch
import numpy as np
import requests
import json
from typing import Optional, Union, List
from pathlib import Path

from neurosync.utils.generate_face_shapes import generate_facial_data_from_bytes
from neurosync.utils.model.model import load_model
from neurosync.utils.config import config
from txt2json import RecordDataJson,Frame
from pipeline.utils.bs61 import *
from dataclasses import dataclass,asdict, field
   
def neurosyncToUnityJson(bs_input,bs_scale = 1.0):            
    rdj = RecordDataJson()
    time = 0
    step = 1.0 / 60
    
    
    
    for i in range(bs_input.shape[0]):
        frame = Frame()
        frame.time = time
        time = time + step
        frame.location = list(range(52))

        weights = []
        sweights = bs_input[i]
        
        for j in range(51):
            weight = bs_input[i,map_52_to_61(j)]
            weights.append((float)(weight) * 100 * bs_scale)
        
        weights.append(0.0)#last tounge bs
        frame.weight = weights
        rdj.frames.append(frame)
    
    
    # jstr = json.dumps(asdict(rdj))
    return rdj
    
class BSGenerator:
    def __init__(self):
        # self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')     
        self.device = torch.device('cpu')   
        self.model_path = os.path.join(os.path.dirname(__file__),config['model_path'])
        self.blendshape_model = load_model(self.model_path, config, self.device)
        
        
    def computeBS(self, wav_bytes):      
        
        start_time = time.time()    
        generated_facial_data = generate_facial_data_from_bytes(wav_bytes, self.blendshape_model, self.device, config)
        bsJson = neurosyncToUnityJson(generated_facial_data)        
        past_time = time.time() - start_time
        print(f"NeuroSync Synthesis finished in {past_time:.4f} sec!")
        return bsJson
         
